import os
import logging
from typing import Optional, Dict, Any, AsyncGenerator
from dotenv import load_dotenv
from camel.configs import QwenConfig
from camel.models import ModelFactory
from camel.types.enums import ModelType, ModelPlatformType
from camel.agents import ChatAgent
from next_question_predictor import init_predictor
import openai

load_dotenv()
DSAPI_KEY = os.getenv('DS_API_KEY')
# 配置日志
logging.basicConfig(level=logging.INFO,
                   format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

VISUALIZATION_PROMPT = """
你是一个专门帮助学生生动形象理解算法和解题思路的助手。你的任务是：

1. 使用贴近生活的比喻和例子
2. 用通俗易懂的语言解释复杂概念
3. 可以使用故事、场景或游戏的方式来解释
4. 结合具体的例子进行讲解（不能是代码）
5. 将抽象概念具象化（不能是代码）

回答要求：
1. 保持简洁明了
2. 避免使用过于专业的术语
3. 多使用类比和比喻
4. 可以加入有趣的元素
5. 注重循序渐进的讲解
6. 禁止使用具体代码进行解释

当用户询问当前题目或未指明其他问题时时（如：解释一下。；解释一下这个题目；帮我理解！）：
1. 分析题目要求和代码逻辑
2. 结合具体的题目内容设计比喻
3. 用生活中的例子解释代码的工作原理

当用户指出分析其他问题时（如：题目为快速排序。用户指出解释希尔排序，你会选择只提及希尔排序）：
1.忽视当前的题目以及编辑器代码
2.专注于分析用户的新问题
3.禁止将话题转回当前题目以及编辑器代码

示例：
问：解释快速排序算法
答：想象你在整理一堆扑克牌：
1. 先随便选一张牌（比如8），这就是我们的"基准牌"
2. 然后把所有牌分成两堆：小于8的放左边，大于8的放右边
3. 对每一堆重复这个过程，直到所有牌都排好

这就像在图书馆整理书架，先按一个标准（如书的厚度）大致分类，然后再细分，最后就能得到整齐的书架。
"""

class VisualizationAgent:
    def __init__(self):
        self.ai_assistant = self._create_ai_assistant()
        self.predictor = None  # 问题预测器
        
        # 添加异步客户端
        self.client = openai.AsyncOpenAI(
            api_key=DSAPI_KEY,
            base_url="https://api.deepseek.com/"
        )

    def _create_ai_assistant(self):
        """创建AI助手实例"""
        try:
            logger.info("正在创建可视化AI助手...")

            
            model = ModelFactory.create(
                model_platform=ModelPlatformType.DEEPSEEK,
                model_type=ModelType.DEEPSEEK_CHAT,
                api_key=DSAPI_KEY,
                url="https://api.deepseek.com"
            )

            agent = ChatAgent(
                system_message=VISUALIZATION_PROMPT,
                model=model,
                message_window_size=10,
                output_language='Chinese'
            )
            logger.info("可视化AI助手创建成功")
            return agent
        except Exception as e:
            logger.error(f"创建可视化AI助手时出错: {str(e)}")
            raise

    async def _stream_chat(self, messages: list) -> AsyncGenerator[str, None]:
        """使用OpenAI API进行流式对话"""
        try:
            stream = await self.client.chat.completions.create(
                model="deepseek-chat",
                messages=messages,
                stream=True
            )
            
            async for chunk in stream:
                if chunk.choices[0].delta.content:
                    yield chunk.choices[0].delta.content
                    
        except Exception as e:
            logger.error(f"流式对话出错: {str(e)}")
            yield f"出错: {str(e)}"

    async def visualize_stream(self, query: str, problem_content: str = "", editor_code: str = "") -> AsyncGenerator[str, None]:
        """生成生动形象的解释（流式版本）"""
        try:
            logger.info(f"开始流式生成可视化解释: {query}")
            logger.info(f"题目内容: {problem_content}")
            logger.info(f"编辑器代码: {editor_code}")
            
            # 构建完整的提示
            full_prompt = query
            if problem_content:
                full_prompt = f"可能被使用也可能不被使用的题目：\n{problem_content}\n\n{query}"
            if editor_code:
                full_prompt = f"{full_prompt}\n\n可能被使用也可能不被使用的当前代码：\n{editor_code}"
            
            # 构建消息
            messages = [
                {"role": "system", "content": VISUALIZATION_PROMPT},
                {"role": "user", "content": full_prompt}
            ]
            
            # 流式生成解释
            explanation = ""
            async for chunk in self._stream_chat(messages):
                explanation += chunk  # 累积完整的解释用于后续预测问题
                yield chunk
                
            # 预测可能的后续问题
            if self.predictor is None:
                logger.info("初始化问题预测器...")
                self.predictor = init_predictor(DSAPI_KEY)
                
            current_context = {
                'problem_content': problem_content,
                'editor_code': editor_code,
                'query': query,
                'visualization_type': 'analogy'
            }
            
            logger.info("开始预测后续问题...")
            next_questions = self.predictor.predict_next_questions(
                current_context=current_context,
                task_response=explanation
            )
            logger.info(f"预测到 {len(next_questions)} 个问题")
            
            # 如果预测器没有返回问题，使用默认问题
            if not next_questions:
                next_questions = [
                    {"question": "能再举一个生活中的例子来解释这个概念吗？"},
                    {"question": "如果遇到更复杂的情况，这个解释还适用吗？"},
                    {"question": "这个类比和实际的代码实现有什么对应关系？"}
                ]
            
            # 在流式输出最后，发送预测问题
            yield "\n\n接下来，您可能想问：\n" + "\n".join([
                f"- {q['question']}" for q in next_questions
            ])
            
        except Exception as e:
            logger.error(f"生成可视化解释时出错: {str(e)}")
            yield f"生成解释失败: {str(e)}"

    def visualize(self, query: str, problem_content: str = "", editor_code: str = "") -> Dict[str, Any]:
        """生成生动形象的解释"""
        try:
            logger.info(f"正在生成可视化解释: {query}")
            logger.info(f"题目内容: {problem_content}")
            logger.info(f"编辑器代码: {editor_code}")
            
            # 构建完整的提示
            full_prompt = query
            if problem_content:
                full_prompt = f"题目：\n{problem_content}\n\n{query}"
            if editor_code:
                full_prompt = f"{full_prompt}\n\n当前代码：\n{editor_code}"
            
            response = self.ai_assistant.step(full_prompt)
            self.ai_assistant.reset()
            
            if not response or not response.msgs:
                logger.warning("模型没有返回任何消息")
                return {
                    'success': False,
                    'response': "生成解释失败，请重试",
                    'predicted_questions': []
                }
                
            explanation = response.msgs[0].content.strip()
            
            # 预测可能的后续问题
            if self.predictor is None:
                logger.info("初始化问题预测器...")
                self.predictor = init_predictor(DSAPI_KEY)
                
            current_context = {
                'problem_content': problem_content,
                'editor_code': editor_code,
                'query': query,
                'visualization_type': 'analogy'
            }
            
            logger.info("开始预测后续问题...")
            next_questions = self.predictor.predict_next_questions(
                current_context=current_context,
                task_response=explanation
            )
            logger.info(f"预测到 {len(next_questions)} 个问题")
            
            # 如果预测器没有返回问题，生成默认的后续问题
            if not next_questions:
                next_questions = [
                    {"question": "能再举一个生活中的例子来解释这个概念吗？"},
                    {"question": "如果遇到更复杂的情况，这个解释还适用吗？"},
                    {"question": "这个类比和实际的代码实现有什么对应关系？"}
                ]
            
            return {
                'success': True,
                'response': explanation,
                'predicted_questions': next_questions
            }
            
        except Exception as e:
            logger.error(f"生成可视化解释时出错: {str(e)}")
            return {
                'success': False,
                'response': f"生成解释失败: {str(e)}",
                'predicted_questions': []
            }
